#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>
LastChangedDate: 2008-03-12 19:52:33 +0000 (Wed, 12 Mar 2008) $

Constructs the reaction information

'''

__revision__ = '$Rev$'

from copy import copy
import os
import libsbml
from swiginac import symbol, numeric

#Import non-standard modules
from VariableDict import Variables


class ReactionInformation(object):
    '''
        Store info about the reaction
    '''
    
    def __init__(self, sbml_id):
        '''Set reaction id'''
        self.sbml_id = sbml_id
        self.stoic = None
        self.rate = None
        self.species_order = None
        
    
    def setStoic(self, stoic):
        '''Store stoichiometry of the reaction'''
        self.stoic = stoic
        return True
    
    def setRate(self, rate):
        '''Store the rate by species coefficent'''
        self.rate = rate
        return True
    
    def setSpeciesOrder(self, species_order):
        '''Stores the species order of the stoic matrix.
           
           XXX inefficient since we store the same info in each
           reaction - oh well.
        '''
        self.species_order = species_order
        return True 




class GetStoichiometryDict(object):
    '''
        Constructs the lambda_{i1, i2} from an sbml file
    '''
        
    def __init__(self, sbml_string=None, sbml_file=None, dir_in=None):
        #Handle the SBML
        r = libsbml.SBMLReader()

        if sbml_string:
            self.d = r.readSBMLFromString(sbml_string)
        else:
            if dir_in == None:
                dir_in = os.getcwd()
            file_path = os.path.join(dir_in, sbml_file)
            assert os.path.exists(file_path), "Invalid path"
            self.d = r.readSBML(file_path)
        self.d.setLevelAndVersion(2, 3)
        self.m = self.d.getModel()
        self.pa_dict = {}
       
        #Create species dictionaires
        self.sp_dict, self.empty_list = self.getSpeciesDict()
        
        #Create parameter dictionaires

        self.updateParameterDict(self.m.getListOfParameters())
        
        self.reaction_list = []

    def getSpeciesDict(self):
        '''
        returns a dictionary of species, of the form:
        sp_name:{'amount':amount, 'index':index}
        empty_list is length(sp_name) vector of zeros
        '''
        sp_dict = {}
        empty_list = []
        for i, sp in enumerate(self.m.getListOfSpecies()):
            if sp.getConstant():
                self.__addParameter(sp.getId(), int(sp.getInitialAmount()))
            else:               
                sp_dict[sp.getId()] = {'amount':int(sp.getInitialAmount())}
                sp_dict[sp.getId()]['index'] = i
                sp_dict[sp.getId()]['name'] = sp.getName()
                empty_list.append(0)
        return sp_dict, empty_list

    def updateParameterDict(self, pa_list):
        '''
        Scans the global parameters and returns a dictionary
        '''
        for i in pa_list:
            if self.pa_dict.has_key(i.getId()):
                raise 'All parameters must be unique'
            else:
                self.__addParameter(i.getId(), i.getValue())
        return True
    
    def __addParameter(self, pa_id, pa_value):
        self.pa_dict[pa_id] = {'value':pa_value, 'symbol': symbol(pa_id)}
        return True
        
    def getReactionList(self):
        '''
            simple getter
        '''
        self.addReactions()
        return self.reaction_list
    
    def addReactions(self):
        '''
            Used to update the delta dict
        '''
        #Get the order of the species
        species_order = copy(self.empty_list)
        for sp, values in self.sp_dict.items():
            species_order[values['index']] = sp
        

        for reaction in self.m.getListOfReactions():
            #Initialise values
            self.lambda_dict = {}
            stoic_values = copy(self.empty_list)
            
            for reactant in reaction.getListOfReactants():
                sp = reactant.getSpecies()
                if self.sp_dict.has_key(sp):
                    r_st = -int(reactant.getStoichiometry())
                    stoic_values[self.sp_dict[sp]['index']] = r_st
                
            for product in reaction.getListOfProducts():
                sp = product.getSpecies()
                if self.sp_dict.has_key(sp):
                    p_st = int(product.getStoichiometry())
                    stoic_values[self.sp_dict[sp]['index']] += p_st
            
            #Convert the kinetic law in swiginac
            kinetic_law = reaction.getKineticLaw()
            self.updateParameterDict(kinetic_law.getListOfParameters())
            formula_list = kinetic_law.getFormula().split()
            swiginac_str = self.string2Swiginac(formula_list)
             

            #Split the rate law in coeff's of species
            self.getSpeciesCoeffs(swiginac_str, self.sp_dict.keys())
            
            #Create a ReactionInformation class
            r = ReactionInformation(reaction.getId())
            r.setStoic(stoic_values)
            r.setRate(self.lambda_dict)
            r.setSpeciesOrder(species_order)
            
            self.reaction_list.append(r)
       

    def getSpeciesCoeffs(self, formula, list_of_species, i1_tup=None):
        
        if not i1_tup:
            i1_tup = copy(self.empty_list)
        sp = list_of_species.pop()
        swiginac_sp = Variables().getParameter(sp)

        for degree in range(formula.ldegree(swiginac_sp), formula.degree(swiginac_sp)+1):
            coeff = formula.expand().coeff(swiginac_sp, degree)
            i1_tup[self.sp_dict[sp]['index']] = degree
            
            if list_of_species:
                self.getSpeciesCoeffs(coeff, copy(list_of_species), i1_tup)
            else:
                self.lambda_dict[tuple(i1_tup)] = self.lambda_dict.get(tuple(i1_tup), 0) + coeff
        return True
 
    def string2Swiginac(self, formula_list, formula=numeric(1)):
        '''
            Bad parser
        '''

        part = formula_list.pop(0)
                
        if part[0] == '(':
            formula_list.insert(0, part[1:])
            formula = (self.string2Swiginac(formula_list, formula))
            formula = self.string2Swiginac(formula_list, formula)
            if not formula_list:
                return formula

        if part[-1] == ')':
            try:
                int(part[:-1])
                return numeric(part[:-1])
            except ValueError:
                return Variables().getParameter(part[:-1])
        if part in self.sp_dict.keys() or part in self.pa_dict.keys():
            if not formula_list:
                return Variables().getParameter(part)
            return self.string2Swiginac(formula_list, Variables().getParameter(part))
            
        if part == '*':
            return formula*self.string2Swiginac(formula_list)
        elif part == '-':
            return formula - self.string2Swiginac(formula_list)
        elif part == '/':
            return formula/self.string2Swiginac(formula_list)
        else:
            return numeric(part)
    
        return formula

        

 
if __name__ == '__main__':
    gdd = GetStoichiometryDict(sbml_file = 'xml_files/tmp.xml')
    lambda_dict = gdd.getReactionList()
    for value in lambda_dict:
        print value.sbml_id, value.species_order, value.stoic
        print value.rate










